Description

This PR improves the backward mirroring implementation. Specifically, it takes into account for each (group of) operator node whether doing backward mirroring can be truly benefitial to the total memory footprint (please refer to test case #1 and #2 below). It also considers the data dependencies between the forward node and its corresponding gradient node. This is because it is possible for the feature maps of a layer to be recomputed without recomputing the layer itself (e.g., the Fully-Connected layer, test case #3). Those improvements allow us to further optimize the memory consumption of our DNN training models.

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at https://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Backward Mirroring Improvements
• Test Case Add some ops #1: RNN Cell
  • In the following graphs, we will be using red arrows to denote backward dependencies (a.k.a. feature maps), which usually are the main contributor of the memory consumption.
  • In the example below that mimics an RNN cell, we shall NOT be doing backward mirroring, because otherwise the total amount of feature maps storage will be doubled.

• Test Case [concurrent-blocking-queue-fix] ConcurrentBlockingQueue::Pop's return… #2: MLP Attention
  • In the example below that mimics the MLP attention (a.k.a. additive attention), we shall be doing backward mirroring, because it can help reduce the total maount of feature maps storage from O(T^2 N) to O(2T N).

• Test Case Update dev branch #3: Fully-Connected Layer
  • In the example below that uses a red node to denote a compute-heavy layer whose gradients are not dependent on the output data entries (e.g., the Fully-Connected layer). The red node can also be put on the mirror path. This enables us to relieve the backward dependency on its feature maps (i.e., inputs) without incurring significant performance overhead.

Comments

FYI, @eric-haibin-lin @szha